"""
下载大模型及向量嵌入模型

hugging face下载dataset时候出现You must be authenticated to access it.问题解决
https://blog.csdn.net/WASEFADG/article/details/134702098

from huggingface_hub import login
login()

token 的获取方式为：huggingface 主页 -> Settings -> Access Tokens -> 新建一个 Read 权限的 token 即可。
"""
import os
os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"
os.environ['HF_HOME'] = '/root/autodl-tmp/cache/'

import platform
from config import config
from pathlib import Path
import modelscope
import huggingface_hub


def download_model(model_name: str) -> str:
    """下载模型"""
    os_name = platform.system().lower()
    root_dir = config['model_dir'][os_name]
    Path(root_dir).mkdir(exist_ok=True)

    model_infos = config['models'][model_name]
    repo_id = model_infos['repo_id']
    local_dir = (Path(root_dir) / model_name).as_posix()

    from_site = model_infos['from']
    if from_site == 'hf-mirror':
        huggingface_hub.snapshot_download(repo_id, local_dir=local_dir,
                                          token="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")
    elif from_site == 'modelscope':
        modelscope.snapshot_download(repo_id, local_dir=local_dir)

    return local_dir